TopkFusion

在指定维度上选取 Top-K 个最大值，并同时输出对应的索引。

该算子在融合实现中同时完成：

  • Top-K 值计算

  • 索引提取

  • 可选的结果排序

其计算过程为： 在指定维度上，对每个 slice 进行排序或选择，输出前 K 个最大元素及其原始索引。

当 sorted = false 时，仅保证输出为 Top-K 元素集合， 不保证其在输出中的顺序；当 sorted = true 时， 输出结果按值从大到小排序。

输入：

  • input - 输入张量的数据地址。 数据类型需与所调用的 TopkFusion 接口类型一致，例如：fp_* 接口对应 float*,``hp_*`` 接口对应 half*,``i16_*`` 接口对应 int16_t*,``i32_*`` 接口对应 int32_t*

  • output - Top-K 值输出张量的数据地址。 数据类型与 input 保持一致。

  • output_index - Top-K 索引输出张量的数据地址，类型固定为 int32_t*。

  • parameter - Top-K 参数结构体指针 TopkParameter*，其定义如下：

    typedef struct TopkParameter {
  // primitive parameter
  OpParameter op_parameter_;
  int k_;                 // 需要选取的 Top-K 个数
  int axis_;              // 执行 Top-K 的维度
  bool sorted_;            // 是否对 Top-K 结果按值排序

  // other parameter
  int dim_size_;           // axis 维度长度
  int outer_loop_num_;     // axis 之前维度展开后的循环次数
  int inner_loop_num_;     // axis 之后维度展开后的循环次数
  void *topk_node_list_;   // 临时 TopK 节点缓冲区
} TopkParameter;

  • core_mask - 核掩码（仅共享存储版本使用）。

输出：

  • output - Top-K 值输出张量的数据地址。

  • output_index - Top-K 索引输出张量的数据地址。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持的数据类型：

    fp32, fp64, int8, int16, int32

  • MT7004 支持的数据类型：

    fp16, fp32, int16, int32

  • output 与 output_index 的布局与输入张量保持一致，仅在 Top-K 维度长度变为 K

  • TopkFusion 算子内部使用 void* 进行类型复用， 但调用时 input 与 output 的实际数据类型 必须与所选接口前缀严格一致。

共享存储版本：

void fp_topk_fusion_s(void *input, void *output, int32_t *output_index, TopkParameter *parameter, int core_mask)
void hp_topk_fusion_s(void *input, void *output, int32_t *output_index, TopkParameter *parameter, int core_mask)
void i16_topk_fusion_s(void *input, void *output, int32_t *output_index, TopkParameter *parameter, int core_mask)
void i32_topk_fusion_s(void *input, void *output, int32_t *output_index, TopkParameter *parameter, int core_mask)

C调用示例：

 1// FT78NE 多核示例
 2#include <stdio.h>
 3#include <topk_fusion.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input = (float *)0xA0000000;
 7    float *output = (float *)0xB0000000;
 8    int32_t *output_index = (int32_t *)0xB1000000;
 9
10    TopkParameter param;
11    param.k_ = 5;
12    param.dim_size_ = 100;
13    param.outer_loop_num_ = 1;
14    param.inner_loop_num_ = 1;
15    param.sorted_ = 1;
16
17    int core_mask = 0xff;
18
19    fp_topk_fusion_s(input, output, output_index, &param, core_mask);
20    return 0;
21}

私有存储版本：

void fp_topk_fusion_p(void *input, void *output, int32_t *output_index, TopkParameter *parameter)
void hp_topk_fusion_p(void *input, void *output, int32_t *output_index, TopkParameter *parameter)
void i16_topk_fusion_p(void *input, void *output, int32_t *output_index, TopkParameter *parameter)
void i32_topk_fusion_p(void *input, void *output, int32_t *output_index, TopkParameter *parameter)

C调用示例：

 1// MT7004 单核示例
 2#include <stdio.h>
 3#include <topk_fusion.h>
 4
 5int main(int argc, char* argv[]) {
 6    half *input = (half *)0x10000000;
 7    half *output = (half *)0x10010000;
 8    int32_t *output_index = (int32_t *)0x10020000;
 9
10    TopkParameter param;
11    param.k_ = 3;
12    param.dim_size_ = 64;
13    param.outer_loop_num_ = 1;
14    param.inner_loop_num_ = 1;
15    param.sorted_ = 0;
16
17    hp_topk_fusion_p(input, output, output_index, &param);
18    return 0;
19}